import os
import json
from paddlex import create_pipeline
from collections import defaultdict


class SealRecognitionProcessor:
    def __init__(self, input_dir="cropped_seals", output_dir="output/seal_recognition",save_res_json=False):
        """
        初始化印章识别处理器

        参数:
            input_dir: 输入目录(包含裁剪后的印章图片)
            output_dir: 输出目录(保存识别结果)
        """
        self.input_dir = input_dir
        self.output_dir = output_dir
        self.pipeline = None
        self.save_res_json = save_res_json

        self.ocr_pipeline = None

        # 创建输出目录
        os.makedirs(self.output_dir, exist_ok=True)

        # 初始化识别流水线
        self._init_pipeline()

    def _init_pipeline(self):
        """初始化识别流水线"""
        print("初始化seal_rec产线...")
        self.pipeline = create_pipeline(pipeline="seal_recognition")
        self.ocr_pipeline = create_pipeline(pipeline="OCR")

    def process_images(self):
        """
        处理裁剪后的印章图片，进行文字识别并输出结构化结果

        返回:
            dict: 识别结果字典
        """
        print("开始执行seal_rec产线处理...")
        result_map = defaultdict(lambda: defaultdict(list))

        # 获取所有图片文件
        image_files = [f for f in os.listdir(self.input_dir)
                       if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

        for img_file in image_files:
            self._process_single_image(img_file, result_map)

        print(f"识别完成，结果已保存到: {self.output_dir}")
        return result_map

    def _process_single_image(self, img_file, result_map):
        print(f"处理文件: {img_file}")
        """处理单个图片文件"""
        # 解析文件名 (格式: "图片名称-文字类型-数字类型-result中该类别的id")
        parts = os.path.splitext(img_file)[0].split('-')
        if len(parts) < 4:
            print(f"跳过文件名格式不正确的文件: {img_file}")
            return

        original_name = parts[0]  # 原始图片名称
        text_type = parts[-3]  # 文字类型
        num_type = parts[-2]  # 数字类型
        class_id = parts[-1]  # 类别ID

        # 完整图片路径
        img_path = os.path.join(self.input_dir, img_file)

        # 进行印章识别
        try:
            num_type = int(num_type)
            output=None
            if num_type in (0,2):

                output = self.pipeline.predict(
                    img_path,
                    use_doc_orientation_classify=False,
                    use_doc_unwarping=False,
                )
            elif num_type in (1,3):
                output = self.ocr_pipeline.predict(
                    img_path,
                    use_doc_orientation_classify=False,
                    use_doc_unwarping=False,
                    use_textline_orientation=False,
                )
            if output:
                self._process_recognition_result(output, img_file, original_name, num_type, class_id, result_map)


            # 保存可视化结果和JSON(可选)
            # for res in output:


        except Exception as e:
            print(f"处理文件 {img_file} 时出错: {str(e)}")
        print("处理完成")

    def _process_recognition_result(self, output, img_file, original_name, num_type, class_id, result_map):
        """处理识别结果"""
        # 公章和人民章分别处理
        if int(num_type) in (0,2):
            for res in output:
                seal_res_list = res.get('seal_res_list', [])

                if seal_res_list:
                    rec_text = seal_res_list[0].get('rec_texts', [])
                    rec_scores = seal_res_list[0].get('rec_scores', [])
                    key = f"{original_name}{os.path.splitext(img_file)[1]}"
                    result_map[key][num_type].append(
                        {class_id: {"rec_text": rec_text, "rec_scores": rec_scores}}
                    )
                # 保存原始json解析结果
                if self.save_res_json:
                    res.save_to_json(self.output_dir)
        if int(num_type) in (1,3):
            for res in output:
                rec_texts= res.get('rec_texts', [])

                if rec_texts:
                    rec_text = rec_texts
                    rec_scores =  res.get('rec_scores', [])
                    key = f"{original_name}{os.path.splitext(img_file)[1]}"
                    result_map[key][num_type].append(
                        {class_id: {"rec_text": rec_text, "rec_scores": rec_scores}}
                    )
                # 保存原始json解析结果
                if self.save_res_json:
                    res.save_to_json(self.output_dir)



if __name__ == "__main__":
    # 使用示例
    processor = SealRecognitionProcessor(
        input_dir="cropped_seals",
        output_dir="output/seal_recognition",
        save_res_json=True
    )
    results = processor.process_images()
